# -*- coding: utf-8 -*-
"""
Created on Wed Jun 13 10:31:30 2018

@author: shenzx
"""

# import module
from os.path import join
from ase import Atoms
import numpy as np
from ase.calculators.abacus import potential_set
from ase.calculators.abacus import basis_set

def potential_list():
    return ['potential_lda', 'potential_pbe', 'potential_pbe_sg15']

def basis_list():
    return ['basis_lda minimal', 'basis_lda medium', 'basis_lda accurate',
            'basis_pbe minimal', 'basis_pbe medium', 'basis_pbe accurate',
            'basis_pbe_sg15 minimal', 'basis_pbe_sg15 medium', 'basis_pbe_sg15 accurate']

def judge_exist_stru(stru = None):
    if stru == None:
        print("-" * 60)
        print("No structure, please set 'stru=atoms' that atoms is a struture from ase")
        print("-" * 60)
        return False
    else:
        return True
        

def read_ase_stru(stru = None, coordinates_type = "Cartesian"):
    if judge_exist_stru(stru):
        atoms_list = []
        atoms_position = []
        atoms_masses = []
        atoms_magnetism = []
        atoms_all = stru.get_chemical_symbols()
        
        # sort atoms according to atoms
        for atoms_all_name in atoms_all:
            temp = True
            for atoms_list_name in atoms_list:
                if  atoms_all_name == atoms_list_name:
                    temp = False
                    break
                        
            if temp:
                atoms_list.append(atoms_all_name)
               
        for atoms_list_name in atoms_list:
            atoms_position.append([])
            atoms_masses.append([])
            atoms_magnetism.append([])
        
        # get position, masses, magnetism from ase atoms                                      
        if coordinates_type == 'Cartesian':
            for i in range(len(atoms_list)):
                for j in range(len(atoms_all)):
                    if atoms_all[j] == atoms_list[i]:
                        atoms_position[i].append(list(stru.get_positions()[j]))
                        atoms_masses[i] = stru.get_masses()[j]
                        atoms_magnetism[i] = list(stru.get_initial_magnetic_moments())
                        
        elif coordinates_type == 'Direct':
            for i in range(len(atoms_list)):
                for j in range(len(atoms_all)):
                    if atoms_all[j]==atoms_list[i]:
                        atoms_position[i].append(list(stru.get_scaled_positions()[j]))
                        atoms_masses[i]=stru.get_masses()[j]
                        atoms_magnetism[i]=list(stru.get_initial_magnetic_moments())
        else:
            raise ValueError("'coordinates_type' is ERROR,"
                                 "please set to 'Cartesian' or 'Direct'")
            
        return atoms_list, atoms_masses, atoms_position, atoms_magnetism 
    
def set_potential(atoms_list = None, potential_path = "./", potential_name = None):
    if atoms_list == None:
        print("-" * 60)
        print("Please set right 'atoms_list' for 'set_potential' function")
        print("-" * 60)
        
    else:
        potential = []
        
        if potential_name == None:
            print("-" * 60)
            print("No inputing 'potential_name',"
                   "we will set it to 'potential_lda'")
            print("-" * 60)
                  
            for atoms_list_name in atoms_list:
                potential.append(join(potential_path, potential_set.potential_lda[atoms_list_name]))
    
        elif type(potential_name) == str:
            print("\n"+"-" * 60)
            print("You set 'potential_name' to %s" % potential_name)
            print("-" * 60)
                  
            if potential_name == 'potential_lda':
                for atoms_list_name in atoms_list:
                    potential.append(join(potential_path, potential_set.potential_lda[atoms_list_name]))
                    
            elif potential_name == 'potential_pbe':
                for atoms_list_name in atoms_list:
                    potential.append(join(potential_path, potential_set.potential_pbe[atoms_list_name]))
                    
            elif potential_name == 'potential_pbe_sg15':
                for atoms_list_name in atoms_list:
                    potential.append(join(potential_path, potential_set.potential_pbe_sg15[atoms_list_name]))
                    
            else:
                raise ValueError("'potential_name' is ERROR,"
                                 "please set to potential_lda,potential_pbe or potential_pbe_sg15")
            
        elif type(potential_name) == list:
            print("\n"+"-" * 60)
            print("You set 'potential_name' to {0}, "
                  "please set 'potential_path' and sure all {1} "
                  "in 'potential_path' ".format(potential_name, potential_name))
            print("-" * 60)
                  
            ele_name = {}
            for i in potential_name:
                with open(join(potential_path, i), 'r') as f:
                    lines = f.readlines()
                
                for line in lines:
                    line = line.replace('=', ' = ')
                    line = line.replace('"', ' " ')
                    data = line.split()
                    if len(data) == 0:
                        continue
                    
                    elif data[0] == 'element':
                        ele_name[data[3]] = i
                        break
                    
                    elif len(data) == 2 and data[1] == 'Element':
                        ele_name[data[0]] = i
                        break
                    
                    else:
                        continue
                                  
            for atoms_list_name in atoms_list:
                potential.append(join(potential_path, ele_name[atoms_list_name]))
    
        else:
            raise ValueError("Please sure what you do!!! "
                             "if 'potential_name=None', we will set it to 'potential_lda'; "
                             "if 'potential_name=str', str should be one of potential_list:"
                             "'potential_lda',"
                             "'potential_pbe',"
                             "'potential_pbe_sg15'; "
                             "if 'potential_name=list', all in list should be in potential_path.")
        
        return potential
    
def set_basis(atoms_list = None, basis_path = "./", basis_name = None):
    if atoms_list == None:
        print("-" * 60)
        print("Please set right 'atoms_list' for 'set_basis' function")
        print("-" * 60)
        
    else:
        basis = []
        
        if basis_name == None:
            print("-" * 60)
            print("No inputing 'basis_name',"
                   "we will set it to 'basis_lda minimal'")
            print("-" * 60)
                  
            for atoms_list_name in atoms_list:
                basis.append(join(basis_path, basis_set.basis_lda['minimal'][atoms_list_name]))
    
        elif type(basis_name) == str:
            print("-" * 60)
            print("You set 'basis_name' to %s" % basis_name)
            print("-" * 60)
            
            #if basis_name == 'basis_lda minimal':
            #    for atoms_list_name in atoms_list:
            #        basis.append(join(basis_path, basis_set.basis_lda['minimal'][atoms_list_name]))
            #        
            #elif basis_name == 'basis_lda medium':
            #    for atoms_list_name in atoms_list:
            #        basis.append(join(basis_path, basis_set.basis_lda['medium'][atoms_list_name]))
            #        
            #elif basis_name == 'basis_lda accurate':
            #    for atoms_list_name in atoms_list:
            #        basis.append(join(basis_path, basis_set.basis_lda['accurate'][atoms_list_name]))
            #        
            #elif basis_name == 'basis_pbe minimal':
            #    for atoms_list_name in atoms_list:
            #        basis.append(join(basis_path, basis_set.basis_pbe['minimal'][atoms_list_name]))
            #        
            #elif basis_name == 'basis_pbe medium':
            #    for atoms_list_name in atoms_list:
            #        basis.append(join(basis_path, basis_set.basis_pbe['medium'][atoms_list_name]))
            #        
            #elif basis_name == 'basis_pbe accurate':
            #    for atoms_list_name in atoms_list:
            #        basis.append(join(basis_path, basis_set.basis_pbe['accurate'][atoms_list_name]))
            #        
            #elif basis_name == 'basis_pbe_sg15 minimal':
            #    for atoms_list_name in atoms_list:
            #        basis.append(join(basis_path, basis_set.basis_pbe_sg15['minimal'][atoms_list_name]))
            #        
            #elif basis_name == 'basis_pbe_sg15 medium':
            #    for atoms_list_name in atoms_list:
            #        basis.append(join(basis_path, basis_set.basis_pbe_sg15['medium'][atoms_list_name]))
            #        
            #elif basis_name == 'basis_pbe_sg15 accurate':
            #    for atoms_list_name in atoms_list:
            #        basis.append(join(basis_path, basis_set.basis_pbe_sg15['accurate'][atoms_list_name]))
            if basis_name in basis_set.basis_dic.keys(): 
                for atoms_list_name in atoms_list:
                    basis.append(join(basis_path, basis_set.basis_dic[basis_name][atoms_list_name]))

            else:
                raise ValueError("'basis_name' is ERROR,"
                                 "please set to one of basis_list："
                                 "'basis_lda minimal',"
                                 "'basis_lda medium',"
                                 "'basis_lda accurate',"
                                 "'basis_pbe minimal',"
                                 "'basis_pbe medium',"
                                 "'basis_pbe accurate',"
                                 "'basis_pbe_sg15 minimal',"
                                 "'basis_pbe_sg15 medium',"
                                 "'basis_pbe_sg15 accurate'")
                            
        elif type(basis_name) == list:
            print("-" * 60)
            print("You set 'basis_name' to {0}, "
                  "please set 'basis_path' and sure all {1} "
                  "in 'basis_path' ".format(basis_name, basis_name))
            print("-" * 60)
                  
            ele_name = {}
            for i in basis_name:
                with open(join(basis_path, i), 'r') as f:
                    lines=f.readlines()
                    
                for line in lines:
                    data = line.split()
                    if len(data) == 0:
                        continue
                    
                    elif data[0] == 'Element':
                        ele_name[data[1]] = i
                        break
                    
                    else:
                        continue
                    
            for atoms_list_name in atoms_list:
                basis.append(join(basis_path, ele_name[atoms_list_name]))
    
        else:
            raise ValueError("Please sure what you do!!! "
                             "if 'basis_name=None', we will set it to basis_lda minimal; "
                             "if 'basis_name=str', str should be one of basis_list: "
                             "'basis_lda minimal',"
                             "'basis_lda medium',"
                             "'basis_lda accurate',"
                             "'basis_pbe minimal',"
                             "'basis_pbe medium',"
                             "'basis_pbe accurate',"
                             "'basis_pbe_sg15 minimal',"
                             "'basis_pbe_sg15 medium',"
                             "'basis_pbe_sg15 accurate';"
                             "if 'potential_name=list', all in list should be in basis_path.")
            
        return basis

def write_input_stru_core(stru = None, directory = "./", filename = "STRU", 
               potential = None, potential_path = "./", basis = None, basis_path = "./", 
               coordinates_type = "Cartesian", 
               atoms_list = None, atoms_position = None, 
               atoms_masses = None, atoms_magnetism = None, fix = 1):
    if not judge_exist_stru(stru):
        return
    
    elif (atoms_list == None or atoms_position == None 
          or atoms_masses == None or atoms_magnetism == None):
        print("-" * 60)
        print("Please set right 'potential', 'basis', "
              "'atoms_list', 'atoms_position', "
              "'atoms_masses', 'atoms_magnetism'")
        print("-" * 60)
        return
        
    else:
        with open(join(directory, filename), 'w') as f :
            f.write('ATOMIC_SPECIES\n')
            for i in range(len(atoms_list)):
                temp1 = ' ' * (4-len(atoms_list[i]))
                temp2 = ' ' * (14-len(str(atoms_masses[i])))
                atomic_species = (atoms_list[i] + temp1
                             + str(atoms_masses[i]) + temp2
                             + potential[i])
                
                f.write(atomic_species)
                f.write('\n')
            
            f.write('\n')
            f.write('NUMERICAL_ORBITAL\n')
            for i in range(len(atoms_list)):
                f.write(basis[i])
                f.write('\n')
            f.write('\n')
            
            f.write('LATTICE_CONSTANT\n')
            f.write('1.889716\n')
            f.write('\n')
            
            f.write('LATTICE_VECTORS\n')
            for i in range(3):
                for j in range(3):
                    temp3 = str("{:0<12f}".format(stru.get_cell()[i][j])) + ' ' * 3
                    f.write(temp3)
                    f.write('   ')
                f.write('\n')
            f.write('\n')
            
            f.write('ATOMIC_POSITIONS\n')
            f.write(coordinates_type)
            f.write('\n')
            f.write('\n')
            for i in range(len(atoms_list)):
                f.write(atoms_list[i])
                f.write('\n')
                f.write(str("{:0<12f}".format(atoms_magnetism[i][0])))
                f.write('\n')
                f.write(str(len(atoms_position[i])))
                f.write('\n')
                
                for j in range(len(atoms_position[i])):
                    temp4 = str("{:0<12f}".format(atoms_position[i][j][0])) + ' ' * 3
                    temp5 = str("{:0<12f}".format(atoms_position[i][j][1])) + ' ' * 3
                    temp6 = str("{:0<12f}".format(atoms_position[i][j][2])) + ' ' * 3
                    sym_pos=(temp4 + temp5 + temp6 +
                             (str(fix) + '   ') * 3)
                    f.write(sym_pos)
                    f.write('\n')
                f.write('\n\n') 
        
        pb_information = {}
        pb_information['potential_path'] = potential_path
        pb_information['basis_path'] = basis_path
        pb_information['potential_name'] = potential
        pb_information['basis_name'] = basis
        return  pb_information
            
def write_input_stru(stru = None, potential_path = './', potential_name = None,
                 basis_path = './', basis_name = None, fix = 1,
                 filename = 'STRU', directory = './',  coordinates_type = 'Cartesian', **kwargs):
    
    if not judge_exist_stru(stru):
        return
    
    else:      
        atoms_list, atoms_masses, atoms_position, atoms_magnetism = read_ase_stru(stru, coordinates_type)
        #print " atoms_list      : ", atoms_list  
        #print " atoms_masses    : ", atoms_masses 
        #print " atoms_magnetism : ", atoms_magnetism 
        #print " atoms_position  : ", atoms_position 

        potential = set_potential(atoms_list, potential_path, potential_name)
        
        basis = set_basis(atoms_list, basis_path, basis_name)
        
        pb_information = write_input_stru_core(stru, directory, filename, 
                                potential, potential_path, basis, basis_path, coordinates_type, 
                                atoms_list, atoms_position, atoms_masses, atoms_magnetism, fix)
        
        return  pb_information

def read_stru(filename='STRU', directory='./', ase=True, **kwargs):
    # Read structure information from abacus structure file
    try:
        f = open(join(directory, filename), 'r')
    except Exception:
        print("-" * 60)
        print("Failed to open 'STRU', Please Check")
        print("-" * 60)
        return 
    else:
        lines = f.readlines()
        f.close()
    
    # initialize reading information
    temp = []
    for line in lines:
        line = line.strip()
        line = line.replace('\n' , ' ')
        line = line.replace('\t' , ' ')
        line = line.replace('//' , ' ')
        line = line.replace('#' , ' ')
                            
        if len(line) != 0:
            temp.append(line)
    
    atom_species = 0
    for i in range(len(temp)):
        if temp[i] == 'NUMERICAL_ORBITAL':
            atom_species = i - 1
            break
        
    atom_symbol = []
    atom_mass = []
    atom_potential = []
    atom_number = []
    atom_magnetism = []
    atom_positions = []
    atom_fix = []
    
    # get symbol, mass, potential
    for i in range(1, atom_species+1):
        atom_symbol.append(temp[i].split()[0])
        atom_mass.append(float(temp[i].split()[1]))
        atom_potential.append(temp[i].split()[2])
        atom_number.append(0)
        atom_magnetism.append(0)
        atom_positions.append([])
        atom_fix.append([])
    
    # get basis
    atom_basis = []
    for i in range(atom_species+2, (atom_species+1) * 2):
        atom_basis.append(temp[i].split()[0])
    
    # get lattice
    atom_lattice_scale = float(temp[(atom_species+1) * 2 + 1].split()[0])
    atom_lattice = np.array([[float(temp[(atom_species+1) * 2 + 3 + i].split()[:3][j]) 
    for j in range(3) ] for i in range(3)])
    
    # get coordinates type
    atom_coor = temp[(atom_species + 1) * 2 + 7].split()[0]
    
    # get position,  atoms number, magnetism, fix
    for i in range(atom_species):
        pos_start = (atom_species + 1) * 2 + 8 + 3 * i
        for j in range(i):
            pos_start += atom_number[j]
        atom_it = atom_symbol.index(temp[pos_start].split()[0])
        atom_magnetism[atom_it] = float(temp[pos_start + 1].split()[0])
        atom_number[atom_it] = int(temp[pos_start + 2].split()[0])

        atom_positions[atom_it] = np.array([[float(temp[pos_start + 3 + i].split()[:3][j])
                      for j in range(3)]for i in range(atom_number[atom_it])])
        
        atom_fix[atom_it] = np.array([[int(temp[pos_start + 3 + i].split()[3:6][j])
                      for j in range(3)]for i in range(atom_number[atom_it])])
        
    # Reset structure information and return results
    formula_symbol = ''
    formula_positions = []
    for i in range(atom_species):
        if atom_number[i] == 1:
            formula_symbol += atom_symbol[i]
            
        else:   
            formula_symbol += atom_symbol[i] + str(atom_number[i])
            
        for j in range(atom_number[i]):
            formula_positions.append(atom_positions[i][j])
    
    formula_cell = atom_lattice * atom_lattice_scale * 0.529177
    
    if ase == True:
        if atom_coor == 'Direct':
            return Atoms(symbols=formula_symbol, cell=formula_cell, scaled_positions=formula_positions)
        
        elif atom_coor == 'Cartesian':
            return Atoms(symbols=formula_symbol,cell=formula_cell,positions=formula_positions)
        
        else :
            raise ValueError("atomic coordinate type is ERROR,please set to Direct or Cartesian")
        
    else :
        return formula_symbol,formula_cell,formula_positions,atom_potential,atom_basis
       
if __name__ == "__main__":
    print("-" * 60)
    print(potential_list())
    print("-" * 60)
    print(basis_list())
    print("-" * 60)
    
    print(write_abacus(stru = None, potential_path = './', potential_name = None,
                 basis_path = './', basis_name = None, fix = 1,
                 filename = 'STRU', directory = './',  coordinates_type = 'Cartesian'))
        
    print(read_stru(filename='STRU', directory='./', ase=True))
